{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ui3LiQ4HdrnA"
      },
      "source": [
        "# Question answering (PyTorch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DSlT-aF3drnG"
      },
      "source": [
        "Install the Transformers, Datasets, and Evaluate libraries to run this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zE2CH8t8drnH"
      },
      "outputs": [],
      "source": [
        "!pip install datasets evaluate transformers[sentencepiece]\n",
        "!pip install accelerate\n",
        "# To run the training on TPU, you will need to uncomment the following line:\n",
        "# !pip install cloud-tpu-client==0.10 torch==1.9.0 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl\n",
        "!apt install git-lfs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IsNwA-FjdrnJ"
      },
      "source": [
        "You will need to setup git, adapt your email and name in the following cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N39_ybazdrnJ"
      },
      "outputs": [],
      "source": [
        "!git config --global user.email \"kindlytree@163.com\"\n",
        "!git config --global user.name \"kindlytree\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ivf7RU6VdrnK"
      },
      "source": [
        "You will also need to be logged in to the Hugging Face Hub. Execute the following and enter your credentials."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B0Xp_LH-drnK"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OxAZ-mrHdrnM"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "raw_datasets = load_dataset(\"squad\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8gKjBT84drnM",
        "outputId": "f5329901-6b06-455c-e053-034d8e1311f2"
      },
      "outputs": [],
      "source": [
        "raw_datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9QBZld0BdrnO",
        "outputId": "f5356d78-f686-40f0-db49-770658ecf19b"
      },
      "outputs": [],
      "source": [
        "print(\"Context: \", raw_datasets[\"train\"][0][\"context\"])\n",
        "print(\"Question: \", raw_datasets[\"train\"][0][\"question\"])\n",
        "print(\"Answer: \", raw_datasets[\"train\"][0][\"answers\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cNH1UcJedrnO",
        "outputId": "02aab1f3-7226-4a5a-97da-5299cdbde09f"
      },
      "outputs": [],
      "source": [
        "raw_datasets[\"train\"].filter(lambda x: len(x[\"answers\"][\"text\"]) != 1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "validation 数据集中，answers 字段通常包含一个 text 的列表，原因主要有:在评估模型时，如果模型的输出答案与 answers['text'] 列表中的任何一个匹配，通常就被认为是正确的。因此，提供多个参考答案可以让模型在评估时更具灵活性，减少因格式或表述差异带来的不必要的惩罚。有时，数据集会在 answers['text'] 列表中包含相同或非常相似的答案。这样的冗余可能是为了增强模型在面对不完美或重复答案时的鲁棒性。\n",
        "变体：不同的人可能会以不同的方式回答相同的问题。例如，对于问题 \"What is the capital of France?\"，可能的正确答案有 \"Paris\", \"The capital of France is Paris\", 或者只是简单的 \"Paris.\"。因此，answers['text'] 中可能会包含多个变体答案，以覆盖所有合理的回答。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mpiu0GINdrnP",
        "outputId": "bd7ee5d9-9814-4c27-b29c-1a2fdcb71af8"
      },
      "outputs": [],
      "source": [
        "print(raw_datasets[\"validation\"][0][\"answers\"])\n",
        "print(raw_datasets[\"validation\"][2][\"answers\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BruMdux9drnP",
        "outputId": "4468531f-1498-4854-f0d4-964d36b4c4ee"
      },
      "outputs": [],
      "source": [
        "print(raw_datasets[\"validation\"][2][\"context\"])\n",
        "print(raw_datasets[\"validation\"][2][\"question\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Js6neWbmdrnP"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "model_checkpoint = \"bert-base-cased\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zky1WDEHdrnQ",
        "outputId": "3b776571-5a41-448b-c671-08b76b4fee7e"
      },
      "outputs": [],
      "source": [
        "tokenizer.is_fast"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pxQ1ykqsdrnQ",
        "outputId": "5fde218f-ffb9-47e8-b992-5b6deecc56d6"
      },
      "outputs": [],
      "source": [
        "context = raw_datasets[\"train\"][0][\"context\"]\n",
        "question = raw_datasets[\"train\"][0][\"question\"]\n",
        "\n",
        "inputs = tokenizer(question, context)\n",
        "tokenizer.decode(inputs[\"input_ids\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "限定了max_length=100，下面运行的结果中，下一个句子的开头50个token和上一句的最后50个是overalap的。\n",
        "truncation=\"only_second\",表明只在第二部分的参数上去进行截断，也即是问题本身不截断"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jkxq7bjjdrnQ",
        "outputId": "dbe4caec-bce5-4b73-c3ba-a8fbd5a7bf12"
      },
      "outputs": [],
      "source": [
        "inputs = tokenizer(\n",
        "    question,\n",
        "    context,\n",
        "    max_length=100,\n",
        "    truncation=\"only_second\",\n",
        "    stride=50,\n",
        "    return_overflowing_tokens=True,\n",
        ")\n",
        "\n",
        "for ids in inputs[\"input_ids\"]:\n",
        "    print(tokenizer.decode(ids))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "长文本问答：\n",
        "\n",
        "当你将长文本分割成多个片段以便模型处理时，可以使用 overflow_to_sample_mapping 确保所有片段的答案都能被正确映射回对应的原始问题文本。\n",
        "批量处理长文本：\n",
        "\n",
        "在处理一批文本数据时，overflow_to_sample_mapping 可以帮助你追踪哪个片段对应哪个原始文本样本，这对于训练和推理任务非常重要。\n",
        "文本处理和分析：\n",
        "\n",
        "使用 overflow_to_sample_mapping 可以帮助在长文本的多个片段中保持对原始文本的理解和连接，特别是在模型输出结果需要与原始输入文本关联时。\n",
        "\n",
        "下面的结果中生成的4个相对的样本的索引都是0，和上面的结果是一致的。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sI7gmF9hdrnQ",
        "outputId": "b795144c-ba44-4f14-8e44-ea756c4e2a54"
      },
      "outputs": [],
      "source": [
        "inputs = tokenizer(\n",
        "    question,\n",
        "    context,\n",
        "    max_length=100,\n",
        "    truncation=\"only_second\",\n",
        "    stride=50,\n",
        "    return_overflowing_tokens=True,\n",
        "    return_offsets_mapping=True,\n",
        ")\n",
        "inputs.keys()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CT1kzpPEdrnR",
        "outputId": "a423d6a8-6fff-4c4b-e822-59fd497e0d55"
      },
      "outputs": [],
      "source": [
        "inputs[\"overflow_to_sample_mapping\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "tokenizer(text_a, text_b, ...) 这种调用方式是专门为需要输入一对文本的任务设计的，如问答（Question-Context）或自然语言推断（Premise-Hypothesis）。\n",
        "Tokenizer会自动将它们格式化成BERT等模型需要的格式，通常是：\n",
        "[CLS] question_tokens [SEP] context_tokens [SEP]  \n",
        "\n",
        "为什么在QA中常用\"only_second\"?: 因为问题通常很短且包含关键信息，我们不希望截断它。而上下文通常很长，截断它是合理的。我们宁愿丢失部分上下文，也不愿丢失问题。当一个上下文因为太长而被切分成多个块（chunk）时，后一个块会包含前一个块末尾的50个token。\n",
        "\n",
        "return_offsets_mapping=True\n",
        "作用: 这是一个非常有用的元数据功能，它为每个token提供了其在原始文本中的起止字符位置。\n",
        "输出格式: inputs中会包含一个名为 offset_mapping 的字段。它的值是一个列表，每个元素对应一个特征。对于每个特征，其值又是一个列表，其中每个元素是一个元组 (start_char, end_char)。\n",
        "例如 (0, 5) 表示这个token对应原始字符串中从索引0到4的字符。\n",
        "对于特殊token（如[CLS], [SEP]），以及不属于某个原始文本（比如问题部分token的offset mapping在上下文部分会是(0,0)）的token，其偏移量为 (0, 0)。  \n",
        "为什么在QA中至关重要?: 模型预测出的答案是start_token_index和end_token_index。我们需要用offset_mapping来将这两个token级别的索引映射回原始上下文字符串中的位置，从而抽取出人类可读的答案文本。  \n",
        "start_char = offset_mapping[start_token_index][0]  \n",
        "end_char = offset_mapping[end_token_index][1]  \n",
        "answer_text = context[start_char : end_char]  \n",
        "\n",
        "为什么需要stride（重叠）?: 为了防止答案恰好出现在两个块的切割处，导致答案被“切断”。通过重叠，可以保证答案完整地出现在至少一个块中。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cxEm5k6sdrnR",
        "outputId": "09e35997-08f1-4884-d457-aae4822f4007"
      },
      "outputs": [],
      "source": [
        "inputs = tokenizer(\n",
        "    raw_datasets[\"train\"][2:6][\"question\"],\n",
        "    raw_datasets[\"train\"][2:6][\"context\"],\n",
        "    max_length=100,\n",
        "    truncation=\"only_second\",\n",
        "    stride=50,\n",
        "    return_overflowing_tokens=True,\n",
        "    return_offsets_mapping=True,\n",
        ")\n",
        "\n",
        "print(f\"The 4 examples gave {len(inputs['input_ids'])} features.\")\n",
        "print(f\"Here is where each comes from: {inputs['overflow_to_sample_mapping']}.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "这次有四个问题，对应于训练集中的索引为2，3，4，5个问答数据，每一个分别产生了4，4，4，7个片段（每个片段和上一个片段有50个token的重叠），不丢失上下文信息，问题不截断。总的下来有19个输入token序列。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "len(inputs['offset_mapping'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "offset_mapping 是 Hugging Face 的 tokenizer 在返回的编码结果中提供的一个参数，它用于指示每个 token 在原始文本中的起始和结束字符索引。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "inputs['offset_mapping']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "offset_mapping是相对于整个question+context的字符串中每个token的start和end的索引,是一个\"二维的列表\"，第一个维度为输入样本的token序列，第二个维度为每个序列中每个token的start和end的位置。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "raw_datasets[\"train\"][2:6][\"answers\"][0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "len(inputs.sequence_ids(0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "sequence_ids 的作用\n",
        "sequence_ids：返回一个列表，其中的每个元素对应一个token在整个输入序列中的位置标识。通常，列表的每个元素值为：\n",
        "0 表示该token属于第一个输入序列（例如 question）。\n",
        "1 表示该token属于第二个输入序列（例如 context）。\n",
        "None 表示该位置没有对应的token（例如填充的部分或分隔符）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JdpUHtk0drnR",
        "outputId": "d70cc1e9-d037-44b0-f80e-b9b1e9dd2565"
      },
      "outputs": [],
      "source": [
        "answers = raw_datasets[\"train\"][2:6][\"answers\"]\n",
        "start_positions = []\n",
        "end_positions = []\n",
        "\n",
        "for i, offset in enumerate(inputs[\"offset_mapping\"]): #i为样本的序列的索引, 这个样本是已经经过截断和overflow处理的从原始的问题上下文中产生的新的样本序列\n",
        "    sample_idx = inputs[\"overflow_to_sample_mapping\"][i]#第i个样本对应的原始的question，context的索引位置\n",
        "    answer = answers[sample_idx]#针对这个question，context所对应样本的问题的答案\n",
        "    start_char = answer[\"answer_start\"][0] #answer_start代表问题在context中的offset\n",
        "    end_char = answer[\"answer_start\"][0] + len(answer[\"text\"][0])#问答的结束的offset\n",
        "    sequence_ids = inputs.sequence_ids(i) #当前样本序列的每一token位置是来自question（0）还是context（1）\n",
        "\n",
        "    # Find the start and end of the context\n",
        "    idx = 0\n",
        "    while sequence_ids[idx] != 1:\n",
        "        idx += 1\n",
        "    context_start = idx\n",
        "    while sequence_ids[idx] == 1:\n",
        "        idx += 1\n",
        "    context_end = idx - 1\n",
        "\n",
        "    # If the answer is not fully inside the context, label is (0, 0)\n",
        "    # 如果上下文开始位置的token序列的第一个字符的位置比start_char要大，也就是答案开始的位置在context之前\n",
        "    # 或者上下文结束位置的token序列的最后一个字符的位置比end_char要小，也就是答案结束的字符位置在context之后\n",
        "    if offset[context_start][0] > start_char or offset[context_end][1] < end_char:\n",
        "        start_positions.append(0)\n",
        "        end_positions.append(0)\n",
        "    else:\n",
        "        # Otherwise it's the start and end token positions \n",
        "        #找出答案的起始和结束位置，不是以offset的字符为单位，而是以token的索引为单位\n",
        "        idx = context_start\n",
        "        while idx <= context_end and offset[idx][0] <= start_char:\n",
        "            idx += 1\n",
        "        start_positions.append(idx - 1)\n",
        "\n",
        "        idx = context_end\n",
        "        while idx >= context_start and offset[idx][1] >= end_char:\n",
        "            idx -= 1\n",
        "        end_positions.append(idx + 1)\n",
        "\n",
        "start_positions, end_positions  #通过这个函数找出了每个生成序列答案在各自的序列中的起始位置，如果答案不完全在context里，则标记为（0，0），位置索引是相对于整个序列而言的。19个feature就有19个答案"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QDmXS0QHdrnS",
        "outputId": "8d2e084f-125d-48f3-d8d8-dc06d67d87b4"
      },
      "outputs": [],
      "source": [
        "idx = 0\n",
        "sample_idx = inputs[\"overflow_to_sample_mapping\"][idx]\n",
        "answer = answers[sample_idx][\"text\"][0]\n",
        "\n",
        "start = start_positions[idx]\n",
        "end = end_positions[idx]\n",
        "labeled_answer = tokenizer.decode(inputs[\"input_ids\"][idx][start : end + 1])\n",
        "\n",
        "print(f\"Theoretical answer: {answer}, labels give: {labeled_answer}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wj6DE7-bdrnS",
        "outputId": "e87a210b-cb4a-4e1c-bd1f-a2951b073f26"
      },
      "outputs": [],
      "source": [
        "idx = 5\n",
        "sample_idx = inputs[\"overflow_to_sample_mapping\"][idx]\n",
        "answer = answers[sample_idx][\"text\"][0]\n",
        "\n",
        "decoded_example = tokenizer.decode(inputs[\"input_ids\"][idx])\n",
        "print(f\"Theoretical answer: {answer}, decoded example: {decoded_example}\")\n",
        "\n",
        "start = start_positions[idx]\n",
        "end = end_positions[idx]\n",
        "labeled_answer = tokenizer.decode(inputs[\"input_ids\"][idx][start : end + 1])\n",
        "\n",
        "print(f\"Theoretical answer: {answer}, labels give: {labeled_answer}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ndFD-M0drnS"
      },
      "outputs": [],
      "source": [
        "max_length = 384\n",
        "stride = 128\n",
        "\n",
        "\n",
        "def preprocess_training_examples(examples):\n",
        "    questions = [q.strip() for q in examples[\"question\"]]\n",
        "    inputs = tokenizer(\n",
        "        questions,\n",
        "        examples[\"context\"],\n",
        "        max_length=max_length,\n",
        "        truncation=\"only_second\",\n",
        "        stride=stride,\n",
        "        return_overflowing_tokens=True,\n",
        "        return_offsets_mapping=True,\n",
        "        padding=\"max_length\",\n",
        "    )\n",
        "\n",
        "    offset_mapping = inputs.pop(\"offset_mapping\") #二维列表，\n",
        "    sample_map = inputs.pop(\"overflow_to_sample_mapping\")\n",
        "    answers = examples[\"answers\"]\n",
        "    start_positions = []\n",
        "    end_positions = []\n",
        "\n",
        "    for i, offset in enumerate(offset_mapping):#二维列表，i表示第几个样本\n",
        "        sample_idx = sample_map[i]#表示当前序列所对应的原始的问答的样本\n",
        "        answer = answers[sample_idx]#找出原始问答答案\n",
        "        start_char = answer[\"answer_start\"][0]\n",
        "        end_char = answer[\"answer_start\"][0] + len(answer[\"text\"][0])\n",
        "        sequence_ids = inputs.sequence_ids(i)\n",
        "\n",
        "        # Find the start and end of the context\n",
        "        idx = 0\n",
        "        while sequence_ids[idx] != 1:\n",
        "            idx += 1\n",
        "        context_start = idx\n",
        "        while sequence_ids[idx] == 1:\n",
        "            idx += 1\n",
        "        context_end = idx - 1\n",
        "\n",
        "        # If the answer is not fully inside the context, label is (0, 0)\n",
        "        if offset[context_start][0] > start_char or offset[context_end][1] < end_char:\n",
        "            start_positions.append(0)\n",
        "            end_positions.append(0)\n",
        "        else:\n",
        "            # Otherwise it's the start and end token positions\n",
        "            idx = context_start\n",
        "            while idx <= context_end and offset[idx][0] <= start_char:\n",
        "                idx += 1\n",
        "            start_positions.append(idx - 1)\n",
        "\n",
        "            idx = context_end\n",
        "            while idx >= context_start and offset[idx][1] >= end_char:\n",
        "                idx -= 1\n",
        "            end_positions.append(idx + 1)\n",
        "\n",
        "    inputs[\"start_positions\"] = start_positions\n",
        "    inputs[\"end_positions\"] = end_positions\n",
        "    return inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "raw_datasets[\"train\"].column_names"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lx-p4c71drnS",
        "outputId": "6bb450de-caba-47dc-dd61-637c72a63f4b"
      },
      "outputs": [],
      "source": [
        "train_dataset = raw_datasets[\"train\"].map(\n",
        "    preprocess_training_examples,\n",
        "    batched=True,\n",
        "    remove_columns=raw_datasets[\"train\"].column_names,\n",
        ")\n",
        "len(raw_datasets[\"train\"]), len(train_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "69vijAC2drnS"
      },
      "outputs": [],
      "source": [
        "def preprocess_validation_examples(examples):\n",
        "    questions = [q.strip() for q in examples[\"question\"]]\n",
        "    inputs = tokenizer(\n",
        "        questions,\n",
        "        examples[\"context\"],\n",
        "        max_length=max_length,\n",
        "        truncation=\"only_second\",\n",
        "        stride=stride,\n",
        "        return_overflowing_tokens=True,\n",
        "        return_offsets_mapping=True,\n",
        "        padding=\"max_length\",\n",
        "    )\n",
        "\n",
        "    sample_map = inputs.pop(\"overflow_to_sample_mapping\")\n",
        "    example_ids = []\n",
        "\n",
        "    for i in range(len(inputs[\"input_ids\"])):\n",
        "        sample_idx = sample_map[i]\n",
        "        example_ids.append(examples[\"id\"][sample_idx])\n",
        "\n",
        "        sequence_ids = inputs.sequence_ids(i)\n",
        "        offset = inputs[\"offset_mapping\"][i]\n",
        "        inputs[\"offset_mapping\"][i] = [\n",
        "            o if sequence_ids[k] == 1 else None for k, o in enumerate(offset) #如果是context，则为（start，end），如为question本身则为None\n",
        "        ]\n",
        "\n",
        "    inputs[\"example_id\"] = example_ids\n",
        "    return inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sJlQmTqSdrnT",
        "outputId": "6dded351-ad8c-4a6b-de59-4f78b93608b0"
      },
      "outputs": [],
      "source": [
        "validation_dataset = raw_datasets[\"validation\"].map(\n",
        "    preprocess_validation_examples,\n",
        "    batched=True,\n",
        "    remove_columns=raw_datasets[\"validation\"].column_names,\n",
        ")\n",
        "len(raw_datasets[\"validation\"]), len(validation_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xKLAP97mdrnT"
      },
      "outputs": [],
      "source": [
        "small_eval_set = raw_datasets[\"validation\"].select(range(100))\n",
        "trained_checkpoint = \"distilbert-base-cased-distilled-squad\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(trained_checkpoint)\n",
        "eval_set = small_eval_set.map(\n",
        "    preprocess_validation_examples,\n",
        "    batched=True,\n",
        "    remove_columns=raw_datasets[\"validation\"].column_names,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WPHB4hVEdrnT"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "AutoModelForQuestionAnswering 及其相关的模型在进行推理时，支持字典形式的输入数据。这种字典通常包含模型所需的不同输入张量，如 input_ids、attention_mask 等。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Klq0TggJdrnT"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import AutoModelForQuestionAnswering\n",
        "\n",
        "eval_set_for_model = eval_set.remove_columns([\"example_id\", \"offset_mapping\"])\n",
        "eval_set_for_model.set_format(\"torch\")\n",
        "\n",
        "print(eval_set_for_model.column_names)\n",
        "\n",
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "batch = {k: eval_set_for_model[k].to(device) for k in eval_set_for_model.column_names}\n",
        "trained_model = AutoModelForQuestionAnswering.from_pretrained(trained_checkpoint).to(\n",
        "    device\n",
        ")\n",
        "\n",
        "with torch.no_grad():\n",
        "    outputs = trained_model(**batch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qpp_uLHVdrnT"
      },
      "outputs": [],
      "source": [
        "start_logits = outputs.start_logits.cpu().numpy()\n",
        "end_logits = outputs.end_logits.cpu().numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sj4uP9sXdrnT"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "\n",
        "example_to_features = collections.defaultdict(list)\n",
        "for idx, feature in enumerate(eval_set):\n",
        "    example_to_features[feature[\"example_id\"]].append(idx)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yeUf-gSPdrnU"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "n_best = 20\n",
        "max_answer_length = 30\n",
        "predicted_answers = []\n",
        "\n",
        "for example in small_eval_set:\n",
        "    example_id = example[\"id\"]\n",
        "    context = example[\"context\"]\n",
        "    answers = []\n",
        "\n",
        "    for feature_index in example_to_features[example_id]:\n",
        "        start_logit = start_logits[feature_index]\n",
        "        end_logit = end_logits[feature_index]\n",
        "        offsets = eval_set[\"offset_mapping\"][feature_index]\n",
        "\n",
        "        start_indexes = np.argsort(start_logit)[-1 : -n_best - 1 : -1].tolist()\n",
        "        end_indexes = np.argsort(end_logit)[-1 : -n_best - 1 : -1].tolist()\n",
        "        for start_index in start_indexes:\n",
        "            for end_index in end_indexes:\n",
        "                # Skip answers that are not fully in the context\n",
        "                if offsets[start_index] is None or offsets[end_index] is None:\n",
        "                    continue\n",
        "                # Skip answers with a length that is either < 0 or > max_answer_length.\n",
        "                if (\n",
        "                    end_index < start_index\n",
        "                    or end_index - start_index + 1 > max_answer_length\n",
        "                ):\n",
        "                    continue\n",
        "\n",
        "                answers.append(\n",
        "                    {\n",
        "                        \"text\": context[offsets[start_index][0] : offsets[end_index][1]],\n",
        "                        \"logit_score\": start_logit[start_index] + end_logit[end_index],\n",
        "                    }\n",
        "                )\n",
        "\n",
        "    best_answer = max(answers, key=lambda x: x[\"logit_score\"])\n",
        "    predicted_answers.append({\"id\": example_id, \"prediction_text\": best_answer[\"text\"]})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FXZu8MczdrnU"
      },
      "outputs": [],
      "source": [
        "import evaluate\n",
        "\n",
        "metric = evaluate.load(\"squad\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CsTmr3wcdrnU"
      },
      "outputs": [],
      "source": [
        "theoretical_answers = [\n",
        "    {\"id\": ex[\"id\"], \"answers\": ex[\"answers\"]} for ex in small_eval_set\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9DavjN4PdrnU",
        "outputId": "5a2b12f0-11d1-4a54-f8dd-4af162dc8451"
      },
      "outputs": [],
      "source": [
        "print(predicted_answers[0])\n",
        "print(theoretical_answers[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H74MlmUUdrnU",
        "outputId": "fb8d2de4-a249-4795-ad98-8abc83a87bb1"
      },
      "outputs": [],
      "source": [
        "metric.compute(predictions=predicted_answers, references=theoretical_answers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "El7C8bGhdrnV"
      },
      "outputs": [],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "\n",
        "def compute_metrics(start_logits, end_logits, features, examples):\n",
        "    example_to_features = collections.defaultdict(list)\n",
        "    for idx, feature in enumerate(features):\n",
        "        example_to_features[feature[\"example_id\"]].append(idx)\n",
        "\n",
        "    predicted_answers = []\n",
        "    for example in tqdm(examples):\n",
        "        example_id = example[\"id\"]\n",
        "        context = example[\"context\"]\n",
        "        answers = []\n",
        "\n",
        "        # Loop through all features associated with that example\n",
        "        for feature_index in example_to_features[example_id]:\n",
        "            start_logit = start_logits[feature_index]\n",
        "            end_logit = end_logits[feature_index]\n",
        "            offsets = features[feature_index][\"offset_mapping\"]\n",
        "\n",
        "            start_indexes = np.argsort(start_logit)[-1 : -n_best - 1 : -1].tolist()\n",
        "            end_indexes = np.argsort(end_logit)[-1 : -n_best - 1 : -1].tolist()\n",
        "            for start_index in start_indexes:\n",
        "                for end_index in end_indexes:\n",
        "                    # Skip answers that are not fully in the context\n",
        "                    if offsets[start_index] is None or offsets[end_index] is None:\n",
        "                        continue\n",
        "                    # Skip answers with a length that is either < 0 or > max_answer_length\n",
        "                    if (\n",
        "                        end_index < start_index\n",
        "                        or end_index - start_index + 1 > max_answer_length\n",
        "                    ):\n",
        "                        continue\n",
        "\n",
        "                    answer = {\n",
        "                        \"text\": context[offsets[start_index][0] : offsets[end_index][1]],\n",
        "                        \"logit_score\": start_logit[start_index] + end_logit[end_index],\n",
        "                    }\n",
        "                    answers.append(answer)\n",
        "\n",
        "        # Select the answer with the best score\n",
        "        if len(answers) > 0:\n",
        "            best_answer = max(answers, key=lambda x: x[\"logit_score\"])\n",
        "            predicted_answers.append(\n",
        "                {\"id\": example_id, \"prediction_text\": best_answer[\"text\"]}\n",
        "            )\n",
        "        else:\n",
        "            predicted_answers.append({\"id\": example_id, \"prediction_text\": \"\"})\n",
        "\n",
        "    theoretical_answers = [{\"id\": ex[\"id\"], \"answers\": ex[\"answers\"]} for ex in examples]\n",
        "    return metric.compute(predictions=predicted_answers, references=theoretical_answers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-K1-kPJAdrnV",
        "outputId": "9195ae3d-b6ae-4a28-dd48-620ac69e06e8"
      },
      "outputs": [],
      "source": [
        "compute_metrics(start_logits, end_logits, eval_set, small_eval_set)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RnZ43WIQdrne"
      },
      "outputs": [],
      "source": [
        "model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xr8zdfXgdrne"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YAQt6MQfdrne"
      },
      "outputs": [],
      "source": [
        "from transformers import TrainingArguments\n",
        "\n",
        "args = TrainingArguments(\n",
        "    \"bert-finetuned-squad\",\n",
        "    evaluation_strategy=\"no\",\n",
        "    save_strategy=\"epoch\",\n",
        "    learning_rate=2e-5,\n",
        "    num_train_epochs=3,\n",
        "    weight_decay=0.01,\n",
        "    fp16=False,\n",
        "    push_to_hub=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s6bmbvBHdrnf"
      },
      "outputs": [],
      "source": [
        "from transformers import Trainer\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=args,\n",
        "    train_dataset=train_dataset,\n",
        "    eval_dataset=validation_dataset,\n",
        "    tokenizer=tokenizer,\n",
        ")\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZgX1ZTY_drnf",
        "outputId": "b6f45deb-eac2-42ab-cede-67816841f077"
      },
      "outputs": [],
      "source": [
        "predictions, _, _ = trainer.predict(validation_dataset)\n",
        "start_logits, end_logits = predictions\n",
        "compute_metrics(start_logits, end_logits, validation_dataset, raw_datasets[\"validation\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6z8q8Vardrnf",
        "outputId": "15d8a2b4-ff0f-48f8-ee82-bd932e584bf5"
      },
      "outputs": [],
      "source": [
        "trainer.push_to_hub(commit_message=\"Training complete\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ApbgPVCWdrnf"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "from transformers import default_data_collator\n",
        "\n",
        "train_dataset.set_format(\"torch\")\n",
        "validation_set = validation_dataset.remove_columns([\"example_id\", \"offset_mapping\"])\n",
        "validation_set.set_format(\"torch\")\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    train_dataset,\n",
        "    shuffle=True,\n",
        "    collate_fn=default_data_collator,\n",
        "    batch_size=8,\n",
        ")\n",
        "eval_dataloader = DataLoader(\n",
        "    validation_set, collate_fn=default_data_collator, batch_size=8\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ymfdp15Bdrnf"
      },
      "outputs": [],
      "source": [
        "model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tWC0tpWrdrnf"
      },
      "outputs": [],
      "source": [
        "from torch.optim import AdamW\n",
        "\n",
        "optimizer = AdamW(model.parameters(), lr=2e-5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W9OsF5rLdrng"
      },
      "outputs": [],
      "source": [
        "from accelerate import Accelerator\n",
        "\n",
        "accelerator = Accelerator(fp16=True)\n",
        "model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(\n",
        "    model, optimizer, train_dataloader, eval_dataloader\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zdvabmGSdrng"
      },
      "outputs": [],
      "source": [
        "from transformers import get_scheduler\n",
        "\n",
        "num_train_epochs = 3\n",
        "num_update_steps_per_epoch = len(train_dataloader)\n",
        "num_training_steps = num_train_epochs * num_update_steps_per_epoch\n",
        "\n",
        "lr_scheduler = get_scheduler(\n",
        "    \"linear\",\n",
        "    optimizer=optimizer,\n",
        "    num_warmup_steps=0,\n",
        "    num_training_steps=num_training_steps,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5odvCZoldrng",
        "outputId": "c5fd1901-a9d5-4ea2-c3d9-9e4e06e1e9fb"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import Repository, get_full_repo_name\n",
        "\n",
        "model_name = \"bert-finetuned-squad-accelerate\"\n",
        "repo_name = get_full_repo_name(model_name)\n",
        "repo_name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I93x6C1cdrng"
      },
      "outputs": [],
      "source": [
        "output_dir = \"bert-finetuned-squad-accelerate\"\n",
        "repo = Repository(output_dir, clone_from=repo_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hVyTrHbodrng"
      },
      "outputs": [],
      "source": [
        "from tqdm.auto import tqdm\n",
        "import torch\n",
        "\n",
        "progress_bar = tqdm(range(num_training_steps))\n",
        "\n",
        "for epoch in range(num_train_epochs):\n",
        "    # Training\n",
        "    model.train()\n",
        "    for step, batch in enumerate(train_dataloader):\n",
        "        outputs = model(**batch)\n",
        "        loss = outputs.loss\n",
        "        accelerator.backward(loss)\n",
        "\n",
        "        optimizer.step()\n",
        "        lr_scheduler.step()\n",
        "        optimizer.zero_grad()\n",
        "        progress_bar.update(1)\n",
        "\n",
        "    # Evaluation\n",
        "    model.eval()\n",
        "    start_logits = []\n",
        "    end_logits = []\n",
        "    accelerator.print(\"Evaluation!\")\n",
        "    for batch in tqdm(eval_dataloader):\n",
        "        with torch.no_grad():\n",
        "            outputs = model(**batch)\n",
        "\n",
        "        start_logits.append(accelerator.gather(outputs.start_logits).cpu().numpy())\n",
        "        end_logits.append(accelerator.gather(outputs.end_logits).cpu().numpy())\n",
        "\n",
        "    start_logits = np.concatenate(start_logits)\n",
        "    end_logits = np.concatenate(end_logits)\n",
        "    start_logits = start_logits[: len(validation_dataset)]\n",
        "    end_logits = end_logits[: len(validation_dataset)]\n",
        "\n",
        "    metrics = compute_metrics(\n",
        "        start_logits, end_logits, validation_dataset, raw_datasets[\"validation\"]\n",
        "    )\n",
        "    print(f\"epoch {epoch}:\", metrics)\n",
        "\n",
        "    # Save and upload\n",
        "    accelerator.wait_for_everyone()\n",
        "    unwrapped_model = accelerator.unwrap_model(model)\n",
        "    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)\n",
        "    if accelerator.is_main_process:\n",
        "        tokenizer.save_pretrained(output_dir)\n",
        "        repo.push_to_hub(\n",
        "            commit_message=f\"Training in progress epoch {epoch}\", blocking=False\n",
        "        )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zzxWN_WXdrng"
      },
      "outputs": [],
      "source": [
        "accelerator.wait_for_everyone()\n",
        "unwrapped_model = accelerator.unwrap_model(model)\n",
        "unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bu3Z6c1hdrnh",
        "outputId": "f27ba3e1-7f2b-4b80-df52-eb908704b231"
      },
      "outputs": [],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "# Replace this with your own checkpoint\n",
        "model_checkpoint = \"huggingface-course/bert-finetuned-squad\"\n",
        "question_answerer = pipeline(\"question-answering\", model=model_checkpoint)\n",
        "\n",
        "context = \"\"\"\n",
        "🤗 Transformers is backed by the three most popular deep learning libraries — Jax, PyTorch and TensorFlow — with a seamless integration\n",
        "between them. It's straightforward to train your models with one before loading them for inference with the other.\n",
        "\"\"\"\n",
        "question = \"Which deep learning libraries back 🤗 Transformers?\"\n",
        "question_answerer(question=question, context=context)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Question answering (PyTorch)",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
